# -*- coding: utf-8 -*-
import pandas as pd
import pymysql
from sqlalchemy import create_engine, text
import os
import warnings
import cryptography
warnings.filterwarnings('ignore')

class CSVToMySQLImporter:
    def __init__(self, host, port, user, password, database):
        """
        初始化数据库连接
        
        参数:
        host: MySQL主机地址
        port: MySQL端口
        user: 用户名
        password: 密码
        database: 数据库名
        """
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.database = database
        self.connection_url = f"mysql+pymysql://{user}:{password}@{host}:{port}/{database}"
        self.engine = None
        self.connect()
    
    def connect(self):
        """建立数据库连接"""
        try:
            self.engine = create_engine(self.connection_url)
            print(f"成功连接到数据库: {self.host}:{self.port}/{self.database}")
        except Exception as e:
            print(f"数据库连接失败: {e}")
            raise
    
    def create_tables(self):
        """根据文档说明创建数据表"""
        
        table_definitions = {
            'stocks': """
                CREATE TABLE IF NOT EXISTS stocks (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    ts_code VARCHAR(20) NOT NULL,
                    symbol VARCHAR(10),
                    name VARCHAR(100),
                    area VARCHAR(50),
                    industry VARCHAR(100),
                    list_date DATE,
                    INDEX idx_ts_code (ts_code)
                )
            """,
            'stock_daily_data': """
                CREATE TABLE IF NOT EXISTS stock_daily_data (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    ts_code VARCHAR(20) NOT NULL,
                    trade_date DATE,
                    open DECIMAL(15, 2),
                    close DECIMAL(15, 2),
                    high DECIMAL(15, 2),
                    low DECIMAL(15, 2),
                    vol DECIMAL(20, 2),
                    amount DECIMAL(20, 3),
                    FOREIGN KEY (ts_code) REFERENCES stocks(ts_code)
                )
            """,
            'stock_technical_indicators_clean': """
                CREATE TABLE IF NOT EXISTS stock_technical_indicators_clean (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    ts_code VARCHAR(20) NOT NULL,
                    trade_date DATE,
                    ma5 DECIMAL(15, 2),
                    ma20 DECIMAL(15, 2),
                    rsi DECIMAL(10, 2),
                    macd DECIMAL(15, 4),
                    boll_upper DECIMAL(15, 2),
                    boll_lower DECIMAL(15, 2)
                )
            """,
            'ai_predictions': """
                CREATE TABLE IF NOT EXISTS ai_predictions (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    ts_code VARCHAR(20) NOT NULL,
                    predict_date DATE,
                    for_date DATE,
                    prediction_score DECIMAL(5, 2)
                )
            """
        }
        
        try:
            with self.engine.connect() as conn:
                for table_name, create_sql in table_definitions.items():
                    conn.execute(text(create_sql))
                    print(f"表 {table_name} 创建成功")
                conn.commit()
            print("所有表创建完成!")
        except Exception as e:
            print(f"创建表时发生错误: {e}")
            raise
    
    def import_csv_data(self, csv_directory='.'):
        """
        导入CSV文件数据到数据库
        
        参数:
        csv_directory: CSV文件所在目录，默认为当前目录
        """
        csv_files = {
            'stocks.csv': 'stocks',
            'stock_daily_data.csv': 'stock_daily_data',
            'stock_technical_indicators_clean.csv': 'stock_technical_indicators_clean',
            'ai_predictions.csv': 'ai_predictions'
        }
        
        imported_tables = []
        
        for csv_file, table_name in csv_files.items():
            csv_path = os.path.join(csv_directory, csv_file)
            
            if not os.path.exists(csv_path):
                print(f"警告: 文件 {csv_file} 不存在，跳过导入")
                continue
            
            try:
                # 读取CSV文件
                print(f"正在读取 {csv_file}...")
                df = pd.read_csv(csv_path)
                
                # 显示数据预览
                print(f"数据预览 - 前3行:")
                print(df.head(3))
                print(f"数据形状: {df.shape}")
                
                # 导入数据到数据库
                print(f"正在导入数据到表 {table_name}...")
                df.to_sql(
                    name=table_name,
                    con=self.engine,
                    if_exists='append',  # 追加模式
                    index=False,         # 不包含索引列
                    method='multi',      # 批量插入
                    chunksize=1000       # 每批1000条记录
                )
                
                print(f"成功导入 {len(df)} 条记录到表 {table_name}")
                imported_tables.append(table_name)
                
            except Exception as e:
                print(f"导入 {csv_file} 时发生错误: {e}")
        
        return imported_tables
    
    def verify_data_import(self):
        """验证数据导入结果"""
        tables = ['stocks', 'stock_daily_data', 'stock_technical_indicators_clean', 'ai_predictions']
        
        print("\n数据导入验证:")
        print("-" * 50)
        
        try:
            with self.engine.connect() as conn:
                for table in tables:
                    result = conn.execute(text(f"SELECT COUNT(*) as count FROM {table}"))
                    count = result.fetchone()[0]
                    print(f"表 {table}: {count} 条记录")
        except Exception as e:
            print(f"验证数据时发生错误: {e}")
    
    def close(self):
        """关闭数据库连接"""
        if self.engine:
            self.engine.dispose()
            print("数据库连接已关闭")

def main():
    """
    主函数 - 您需要在这里配置您的数据库连接信息
    """
    # =============================================
    # 配置数据库连接信息
    # =============================================
    db_config = {
        'host': 'mysql-10fed60c-project-manonghezuoshe.g.aivencloud.com',      # 我们要用的MySQL主机地址
        'port': 23966,            # 我们要用的MySQL端口
        'user': 'avnadmin', # 我们要用的用户名
        'password': '', # 我们要用的密码（小组内部保密）
        'database': 'defaultdb'   # 我们要用的数据库名
    }
    
    # =============================================
    # 程序执行部分
    # =============================================
    
    # 创建导入器实例
    importer = CSVToMySQLImporter(**db_config)
    
    try:
        # 步骤1: 创建数据表
        print("步骤1: 创建数据表...")
        importer.create_tables()
        
        # 步骤2: 导入CSV数据
        print("\n步骤2: 导入CSV数据...")
        imported_tables = importer.import_csv_data()
        
        # 步骤3: 验证导入结果
        print("\n步骤3: 验证导入结果...")
        importer.verify_data_import()
        
        print(f"\n导入完成! 成功导入 {len(imported_tables)} 个表的数据")
        
    except Exception as e:
        print(f"程序执行过程中发生错误: {e}")
    finally:
        # 关闭连接
        importer.close()

if __name__ == "__main__":
    main()